package worker

import (
	"context"
	"gitee.com/taotao_today/crontab/crontab/common"
	clientv3 "go.etcd.io/etcd/client/v3"
)

// JobLock 分布式锁(TXN事务)
type JobLock struct {
	kv    clientv3.KV
	lease clientv3.Lease

	// 每一个任务对应一把分布式锁
	jobName string

	cancelFunc context.CancelFunc // 用于终止自动续租
	leaseId    clientv3.LeaseID   // 租约ID
	isLocked   bool               // 是否上锁成功
}

func InitJobLock(jobName string, kv clientv3.KV, lease clientv3.Lease) *JobLock {
	return &JobLock{
		kv:      kv,
		lease:   lease,
		jobName: jobName,
	}
}

func (l *JobLock) TryLock() (err error) {

	// 尝试去ETCD中申请锁

	// 1. 创建租约(5秒)
	leaseGrantResp, err := l.lease.Grant(context.TODO(), 5)
	if err != nil {
		return
	}
	// context用于取消自动续租
	cancelCtx, cancelFunc := context.WithCancel(context.TODO())

	// 租约ID
	leaseId := leaseGrantResp.ID

	// 2. 自动续租
	keepRespChan, err := l.lease.KeepAlive(cancelCtx, leaseId)
	if err != nil {
		l.releaseLease(cancelFunc, leaseId)
		return
	}
	// 3, 处理续租应答的协程
	go func() {
		for {
			select {
			case keepResp := <-keepRespChan:
				if keepResp == nil {
					goto END
				}
			}
		}
	END:
	}()

	// 4. 创建事务
	txn := l.kv.Txn(context.TODO())

	// 锁路径
	lockKey := common.JOB_LOCK_DIR + l.jobName

	// 5, 事务抢锁
	txn.If(clientv3.Compare(clientv3.CreateRevision(lockKey), "=", 0)).
		Then(clientv3.OpPut(lockKey, "", clientv3.WithLease(leaseId))).
		Else(clientv3.OpGet(lockKey))

	// 提交事务
	txnResp, err := txn.Commit()
	if err != nil {
		l.releaseLease(cancelFunc, leaseId)
		return
	}

	// 6, 成功返回, 失败释放租约
	if !txnResp.Succeeded {
		// 锁被占用
		err = common.ERR_LOCK_ALREADY_REQUIRED
		l.releaseLease(cancelFunc, leaseId)
		return
	}

	// 抢锁成功
	l.cancelFunc = cancelFunc
	l.leaseId = leaseId
	l.isLocked = true
	return
}

func (l *JobLock) releaseLease(cancelFunc context.CancelFunc, leaseId clientv3.LeaseID) {
	cancelFunc()
	l.lease.Revoke(context.TODO(), leaseId)
}

func (l *JobLock) Unlock() {
	// 释放锁
	if l.isLocked {
		// 取消我们程序自动续租的协程
		l.cancelFunc()
		// 释放租约
		l.lease.Revoke(context.TODO(), l.leaseId)
	}
}
